// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>

#include <nebula/compute/light_array_internal.h>
#include <nebula/compute/row/encode_internal.h>
#include <nebula/compute/row/row_internal.h>
#include <nebula/compute/util.h>
#include <nebula/core/memory_pool.h>

#include <turbo/utility/status.h>

namespace nebula::compute {

    class TURBO_EXPORT KeyCompare {
    public:
        // Clarify the max temp stack usage for CompareColumnsToRows, which might be necessary
        // for the caller to be aware of (possibly at compile time) to reserve enough stack size
        // in advance. The CompareColumnsToRows implementation uses three uint8 temp vectors as
        // buffers for match vectors, all are of size num_rows. Plus extra kMiniBatchLength to
        // cope with stack padding and aligning.
        constexpr static int64_t CompareColumnsToRowsTempStackUsage(int64_t num_rows) {
            return (sizeof(uint8_t) + sizeof(uint8_t) + sizeof(uint8_t)) * num_rows +
                   /*extra=*/util::MiniBatch::kMiniBatchLength;
        }

        // Returns a single 16-bit selection vector of rows that failed comparison.
        // If there is input selection on the left, the resulting selection is a filtered image
        // of input selection.
        static void CompareColumnsToRows(
                uint32_t num_rows_to_compare, const uint16_t *sel_left_maybe_null,
                const uint32_t *left_to_right_map, LightContext *ctx, uint32_t *out_num_rows,
                uint16_t *out_sel_left_maybe_same, const std::vector<KeyColumnArray> &cols,
                const RowTableImpl &rows, bool are_cols_in_encoding_order,
                uint8_t *out_match_bitvector_maybe_null = nullptr);

    private:
        static uint32_t ColIdInEncodingOrder(const RowTableImpl &rows, uint32_t id_col,
                                             bool are_cols_in_encoding_order) {
            return are_cols_in_encoding_order ? id_col
                                              : rows.metadata().pos_after_encoding(id_col);
        }

        template<bool use_selection>
        static void NullUpdateColumnToRow(uint32_t id_col, uint32_t num_rows_to_compare,
                                          const uint16_t *sel_left_maybe_null,
                                          const uint32_t *left_to_right_map, LightContext *ctx,
                                          const KeyColumnArray &col, const RowTableImpl &rows,
                                          bool are_cols_in_encoding_order,
                                          uint8_t *match_bytevector);

        template<bool use_selection, class COMPARE_FN>
        static void CompareBinaryColumnToRowHelper(
                uint32_t offset_within_row, uint32_t first_row_to_compare,
                uint32_t num_rows_to_compare, const uint16_t *sel_left_maybe_null,
                const uint32_t *left_to_right_map, LightContext *ctx, const KeyColumnArray &col,
                const RowTableImpl &rows, uint8_t *match_bytevector, COMPARE_FN compare_fn);

        template<bool use_selection>
        static void CompareBinaryColumnToRow(uint32_t offset_within_row,
                                             uint32_t num_rows_to_compare,
                                             const uint16_t *sel_left_maybe_null,
                                             const uint32_t *left_to_right_map,
                                             LightContext *ctx, const KeyColumnArray &col,
                                             const RowTableImpl &rows,
                                             uint8_t *match_bytevector);

        template<bool use_selection, bool is_first_varbinary_col>
        static void CompareVarBinaryColumnToRowHelper(
                uint32_t id_varlen_col, uint32_t first_row_to_compare, uint32_t num_rows_to_compare,
                const uint16_t *sel_left_maybe_null, const uint32_t *left_to_right_map,
                LightContext *ctx, const KeyColumnArray &col, const RowTableImpl &rows,
                uint8_t *match_bytevector);

        template<bool use_selection, bool is_first_varbinary_col>
        static void CompareVarBinaryColumnToRow(uint32_t id_varlen_col,
                                                uint32_t num_rows_to_compare,
                                                const uint16_t *sel_left_maybe_null,
                                                const uint32_t *left_to_right_map,
                                                LightContext *ctx, const KeyColumnArray &col,
                                                const RowTableImpl &rows,
                                                uint8_t *match_bytevector);

        static void AndByteVectors(LightContext *ctx, uint32_t num_elements,
                                   uint8_t *bytevector_A, const uint8_t *bytevector_B);

#if NEBULA_HAVE_RUNTIME_AVX2_SUPPORTED

        template<bool use_selection>
        static uint32_t NullUpdateColumnToRowImp_avx2(
                uint32_t id_col, uint32_t num_rows_to_compare, const uint16_t *sel_left_maybe_null,
                const uint32_t *left_to_right_map, LightContext *ctx, const KeyColumnArray &col,
                const RowTableImpl &rows, bool are_cols_in_encoding_order,
                uint8_t *match_bytevector);

        template<bool use_selection, class COMPARE8_FN>
        static uint32_t CompareBinaryColumnToRowHelper_avx2(
                uint32_t offset_within_row, uint32_t num_rows_to_compare,
                const uint16_t *sel_left_maybe_null, const uint32_t *left_to_right_map,
                LightContext *ctx, const KeyColumnArray &col, const RowTableImpl &rows,
                uint8_t *match_bytevector, COMPARE8_FN compare8_fn);

        template<bool use_selection>
        static uint32_t CompareBinaryColumnToRowImp_avx2(
                uint32_t offset_within_row, uint32_t num_rows_to_compare,
                const uint16_t *sel_left_maybe_null, const uint32_t *left_to_right_map,
                LightContext *ctx, const KeyColumnArray &col, const RowTableImpl &rows,
                uint8_t *match_bytevector);

        template<bool use_selection, bool is_first_varbinary_col>
        static void CompareVarBinaryColumnToRowImp_avx2(
                uint32_t id_varlen_col, uint32_t num_rows_to_compare,
                const uint16_t *sel_left_maybe_null, const uint32_t *left_to_right_map,
                LightContext *ctx, const KeyColumnArray &col, const RowTableImpl &rows,
                uint8_t *match_bytevector);

        static uint32_t AndByteVectors_avx2(uint32_t num_elements, uint8_t *bytevector_A,
                                            const uint8_t *bytevector_B);

        static uint32_t NullUpdateColumnToRow_avx2(
                bool use_selection, uint32_t id_col, uint32_t num_rows_to_compare,
                const uint16_t *sel_left_maybe_null, const uint32_t *left_to_right_map,
                LightContext *ctx, const KeyColumnArray &col, const RowTableImpl &rows,
                bool are_cols_in_encoding_order, uint8_t *match_bytevector);

        static uint32_t CompareBinaryColumnToRow_avx2(
                bool use_selection, uint32_t offset_within_row, uint32_t num_rows_to_compare,
                const uint16_t *sel_left_maybe_null, const uint32_t *left_to_right_map,
                LightContext *ctx, const KeyColumnArray &col, const RowTableImpl &rows,
                uint8_t *match_bytevector);

        static uint32_t CompareVarBinaryColumnToRow_avx2(
                bool use_selection, bool is_first_varbinary_col, uint32_t id_varlen_col,
                uint32_t num_rows_to_compare, const uint16_t *sel_left_maybe_null,
                const uint32_t *left_to_right_map, LightContext *ctx, const KeyColumnArray &col,
                const RowTableImpl &rows, uint8_t *match_bytevector);

#endif
    };

}  // namespace nebula::compute
